# -*- coding: utf-8 -*-

import pymysql
import datetime
import re
import time
import os
import traceback
from elasticsearch import helpers
from elasticsearch import Elasticsearch

# 批量从 MySQL 导入数据到 elasticsearch

ES_HOST = "192.168.1.5:9200"
YOUR_DATABASE = "news_spider"
YOUR_TABLE = "tb_test_user"
# # # # # # # # # # # # # # # # # # # # # #

YOUR_DATABASE_USER = "root"
YOUR_DATABASE_PASSWORD = "root"
YOUR_DATABASE_HOST = "127.0.0.1"
YOUR_DATABASE_PORT = 3306


def get_local_mysql():
    try:
        mysql_conn = pymysql.connect(
            host=YOUR_DATABASE_HOST,
            port=YOUR_DATABASE_PORT,
            user=YOUR_DATABASE_USER,
            password=YOUR_DATABASE_PASSWORD,
            db=YOUR_DATABASE,
            charset='utf8'
        )
        return mysql_conn
    except Exception as e:
        print(e)




es = Elasticsearch(hosts=ES_HOST)

settings = {
    "elsticsearch_nodes": [
        ES_HOST
    ]
}
es_client = None
region_dict = {}
es_dict = {}


def get_es_client():
    pid = os.getpid()
    if pid in es_dict:
        return es_dict[pid]
    else:
        es_client = Elasticsearch(
            settings['my-es'],
            sniff_on_start=True,
            sniff_on_connection_fail=True,
            sniffer_timeout=60,
            maxsize=10
        )
        es_dict[pid] = es_client
        return es_client


def es_main():
    mysql_conn = get_local_mysql()
    cursor0 = mysql_conn.cursor()
    cursor0.execute(
        'SELECT * FROM  %s ' % (YOUR_TABLE)
    )
    res0 = cursor0.fetchall()

    list1 = []
    for row in cursor0.description:
        list1.append(row[0])
    tuple(list1)

    actions = []
    tb_name = YOUR_TABLE
    cnt = 0
    for r in res0:
        cnt += 1
        data = dict(zip(list1, r))
        print(cnt)

        try:
            action = {"_index": tb_name, '_id': data["id"], "_source": data}
            actions.append(action)

            # 每1000个一组批量导入
            if len(actions) == 1000:
                time.sleep(1)
                print("execute 1000 action")
                success, errors = helpers.bulk(es, actions, raise_on_error=True)
                actions = []
        except:
            print('traceback.format_exc():\n%s' % traceback.format_exc())
    if len(actions) > 0:
        try:
            success, errors = helpers.bulk(es, actions, raise_on_error=True)
        except Exception as e:
            print(e)


if __name__ == '__main__':
    es_main()